import mmap
import re
from typing import IO, Dict, Optional, Iterator, TypedDict, NewType


Task = NewType("Task", str)
PID = NewType("PID", int)
TracePoints = Dict[PID, Dict[Task, "TracePoint"]]


TRACE_MARK_TAG_REGEX = r"\s*(.+)-(\d+)\s+\(\s*(\d+)\)\s*\[\d+\][\s\.]+(\d+\.\d+):\stracing_mark_write:\s([A-Z])\|(\d+)\|(<tag>.+)\n"
TRACE_MARK_NAME_REGEX = r"\s*(.+)-(\d+)\s+\(\s*(\d+)\)\s*\[\d+\][\s\.]+(\d+\.\d+):\stracing_mark_write:\s([A-Z])\|(\d+)\|(<name>)\n"


class TracePointIterator:
    """迭代trace point"""

    def __init__(self, fp: IO, mm: mmap.mmap = None, start_pos:int=None, end_pos:int=None, tag: Optional[str] = None, name_regex: Optional[str] = None, pattern: Optional[str] = None):
        self.fp = fp
        self.mm = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ) if mm is None else mm
        self.start_pos = start_pos if start_pos is not None else 0
        self.end_pos = end_pos if end_pos is not None else self.mm.size()
        self.tag = tag
        self.name_regex = name_regex
        self._pattern = pattern

    @property
    def pattern(self) -> re.Pattern:
        if not self._pattern:
            if self.tag:
                self._pattern = TRACE_MARK_TAG_REGEX.replace("<tag>", self.tag)
            elif self.name_regex:
                self._pattern = TRACE_MARK_NAME_REGEX.replace("<name>", self.name_regex)
            else:
                raise ValueError("Either tag or name_regex must be provided.")
            self._pattern = re.compile(self._pattern.encode("utf-8"))
        return self._pattern

    def __iter__(self) -> Iterator["TracePoint"]:
        start_pos = self.start_pos
        while True:
            m = self.pattern.search(self.mm, start_pos, self.end_pos)
            if not m:
                break
            _pos_start, _pos_end = m.span()
            tp = TracePoint(pos=_pos_start, mm=self.mm, content=self.mm[_pos_start:_pos_end].decode("utf-8"), size=_pos_end - _pos_start)
            tp.build(m)
            yield tp
            start_pos = _pos_end  # 更新起始位置，继续查找下一个trace point


class TracePoint:
    """单个trace point"""

    def __init__(self, pos: int, mm: mmap.mmap, content:str, size: int):
        if mm is None and fp is None:
            raise ValueError("Either file pointer or memory map must be provided.")
        self.pos:int = pos
        self.size:int = size
        self.mm: mmap.mmap = mm
        self._content:str = content

    def build(self, m: re.Match):
        task, pid, tgid, timestamp, mark, _, name = [_.decode("utf-8") for _ in m.groups()]
        self.task:str = task
        self.pid:int = int(pid)
        self.tgid:int = int(tgid)
        self.timestamp:float = float(timestamp)
        self.mark:str = mark
        self.name:str = name
        
    @property
    def content(self) -> str:
        if not self._content:
            self._content = self.mm[self.pos:self.pos + self.size].decode("utf-8")
        return self._content
    
    def match(self, other: "TracePoint") -> bool:
        """判断两个trace point是否匹配
        B和E匹配(同步任务)，S和F匹配(异步任务)
        """
        return (
            self.task == other.task and
            self.pid == other.pid and
            self.tgid == other.tgid and
            {self.mark, other.mark} in [{"B", "E"}, {"S", "F"}]
        )
    
    def __repr__(self):
        return f"TracePoint(task={self.task}, pid={self.pid}, tgid={self.tgid}, timestamp={self.timestamp}, mark={self.mark}, name={self.name})"


class Span:
    """两个trace point之间的区间"""

    def __init__(self, start_tp: TracePoint = None, end_tp: TracePoint = None, fp: IO = None, mm: mmap.mmap = None):
        if mm is None and fp is None:
            raise ValueError("Either file pointer or memory map must be provided.")
        self.start_tp: TracePoint = start_tp
        self.end_tp: TracePoint = end_tp
        self.mm: mmap.mmap = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ) if mm is None else mm
    
    @property
    def start(self) -> int:
        if not self.start_tp:
            raise ValueError("Start TracePoint is not set.")
        return self.start_tp.pos + self.start_tp.size
    
    @property
    def end(self) -> int:
        if not self.end_tp:
            raise ValueError("End TracePoint is not set.")
        return self.end_tp.pos

    def spans(self, tag: Optional[str]=None, name_regex: Optional[str]=None) -> Iterator["Span"]:
        """在当前span中查找子span"""
        return SpanIterator(fp=None, mm=self.mm, start_pos=self.start, end_pos=self.end, tag=tag, name_regex=name_regex)

    def trace_points(self, tag: Optional[str]=None, name_regex: Optional[str]=None) -> Iterator["TracePoint"]:
        """在当前span中查找trace point"""
        return TracePointIterator(fp=None, mm=self.mm, start_pos=self.start, end_pos=self.end, tag=tag, name_regex=name_regex)

    def get_duration(self) -> float:
        return self.end_tp.timestamp - self.start_tp.timestamp
    
    def __repr__(self):
        return f"Span(start={self.start_tp}, end={self.end_tp})"


class SpanIterator:
    """用正则遍历文件内容"""

    def __init__(self, fp: IO, mm: mmap.mmap = None, start_pos: int = 0, end_pos: Optional[int] = None, tag: Optional[str] = None, name_regex: Optional[str] = None, pattern: Optional[str] = None):
        if mm is None and fp is None:
            raise ValueError("Either file pointer or memory map must be provided.")
        self.fp = fp
        self.mm = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ) if mm is None else mm
        self.start_pos = start_pos
        self.end_pos = end_pos if end_pos is not None else self.mm.size()
        self.tag = tag
        self.name_regex = name_regex
        self._pattern = pattern
        # self.tps: 存储当前span中的trace points，字典形式存储，键为pid，值为TracePoint对象
        self.tps: TracePoints = dict()
    
    @property
    def pattern(self) -> re.Pattern:
        if not self._pattern:
            if self.tag:
                self._pattern = TRACE_MARK_TAG_REGEX.replace("<tag>", self.tag)
            elif self.name_regex:
                self._pattern = TRACE_MARK_NAME_REGEX.replace("<name>", self.name_regex)
            else:
                raise ValueError("Either tag or name_regex must be provided.")
            self._pattern = re.compile(self._pattern.encode("utf-8"))
        return self._pattern

    def __iter__(self) -> Iterator[Span]:
        for tp in TracePointIterator(fp=self.fp, mm=self.mm, start_pos=self.start_pos, end_pos=self.end_pos, pattern=self.pattern):
            if not (self.tps.get(tp.pid) and tp.task in self.tps[tp.pid]):
                self.tps[tp.pid] = {tp.task: tp}
                continue
            # 如果当前tp与上一个tp匹配，则合并span
            start_tp = self.tps[tp.pid][tp.task]
            if start_tp.match(tp):
                del self.tps[tp.pid][tp.task]
                if not self.tps[tp.pid]:
                    del self.tps[tp.pid]
                yield Span(start_tp=start_tp, end_tp=tp, mm=self.mm)
            else:
                print(f"Unmatched TracePoint: {tp.content}")


# 体检规则
with open("test.trace", "rb") as fp:
    for app_list_fling_span in SpanIterator(fp, tag="H:APP_LIST_FLING"):
        print(app_list_fling_span.get_duration())
        print(app_list_fling_span)
        # 滑动操作
        for refresh_rate_tp in app_list_fling_span.trace_points(tag="H:RSHardwareThread::CommitAndReleaseLayers"):
            print(refresh_rate_tp)
            pass
        for waiting_presence_tp in app_list_fling_span.trace_points(tag="H:Waiting for Present Fence"):
            # 获得 等待状态
            print(waiting_presence_tp)



"""
找区域 span
找打点 tracePoint

区域底下会有子区域跟子打点

先找到页面滑动的区域 (H:APP_LIST_FLING)
然后在这个区域下找打点，
第一类打点是系统设置的刷新帧率 (H:RSHardwareThread::CommitAndReleaseLayers)
第二类打点是实际刷新帧率 (H:Waiting for Present Fence)
"""